#include "ompl/base/goals/GoalSampleableRegion.h"
#include "pRRTIterationCount.h"
#include <boost/thread.hpp>

using namespace ompl;

pRRTIterationCount::pRRTIterationCount(const ompl::base::SpaceInformationPtr &si, int numThreads)
	: ompl::geometric::pRRT(si)
{
	setThreadCount(numThreads);
}

pRRTIterationCount::~pRRTIterationCount(void)
{
}

void pRRTIterationCount::clear(void)
{
	ompl::geometric::pRRT::clear();
	uiNumIterations_ = 0;
}

// ++uiNumIterations_ is added
ompl::base::PlannerStatus pRRTIterationCount::solve(const ompl::base::PlannerTerminationCondition &ptc)
{
	base::GoalRegion *goal = dynamic_cast<base::GoalRegion*>(pdef_->getGoal().get());

	if (!goal)
	{
		OMPL_ERROR("Goal undefined");
		return base::PlannerStatus::UNRECOGNIZED_GOAL_TYPE;
	}

	samplerArray_.resize(threadCount_);

	while (const base::State *st = pis_.nextStart())
	{
		Motion *motion = new Motion(si_);
		si_->copyState(motion->state, st);
		nn_->add(motion);
	}

	if (nn_->size() == 0)
	{
		OMPL_ERROR("There are no valid initial states!");
		return base::PlannerStatus::INVALID_START;
	}

	OMPL_INFORM("Starting with %u states", nn_->size());

	SolutionInfoIterationCount sol;
	sol.solution = NULL;
	sol.approxsol = NULL;
	sol.approxdif = std::numeric_limits<double>::infinity();

	sol.iterations.resize(threadCount_);
	for (unsigned int i = 0; i < threadCount_; ++i)
		sol.iterations[i] = 0;

	std::vector<boost::thread*> th(threadCount_);
	for (unsigned int i = 0 ; i < threadCount_ ; ++i)
		th[i] = new boost::thread(boost::bind(&pRRTIterationCount::threadSolveIterationCount, this, i, ptc, &sol));
	for (unsigned int i = 0 ; i < threadCount_ ; ++i)
	{
		th[i]->join();
		delete th[i];
	}

	uiNumIterations_ = sol.iterations[0];
	for (unsigned int i = 1 ; i < threadCount_ ; ++i)
	{
		if (sol.iterations[i] > uiNumIterations_)
			uiNumIterations_ = sol.iterations[i];
	}

	bool solved = false;
	bool approximate = false;
	if (sol.solution == NULL)
	{
		sol.solution = sol.approxsol;
		approximate = true;
	}

	if (sol.solution != NULL)
	{
		lastGoalMotion_ = sol.solution;

		/* construct the solution path */
		std::vector<Motion*> mpath;
		while (sol.solution != NULL)
		{
			mpath.push_back(sol.solution);
			sol.solution = sol.solution->parent;
		}

		/* set the solution path */
		ompl::geometric::PathGeometric *path = new ompl::geometric::PathGeometric(si_);
		for (int i = mpath.size() - 1 ; i >= 0 ; --i)
			path->append(mpath[i]->state);

		pdef_->addSolutionPath(base::PathPtr(path), approximate, sol.approxdif);
		solved = true;
	}

	OMPL_INFORM("Created %u states", nn_->size());

	return base::PlannerStatus(solved, approximate);
}

void pRRTIterationCount::threadSolveIterationCount(unsigned int tid, const ompl::base::PlannerTerminationCondition &ptc, SolutionInfoIterationCount *sol)
{
	checkValidity();
	base::Goal                 *goal   = pdef_->getGoal().get();
	base::GoalSampleableRegion *goal_s = dynamic_cast<base::GoalSampleableRegion*>(goal);
	RNG                         rng;

	Motion *rmotion   = new Motion(si_);
	base::State *rstate = rmotion->state;
	base::State *xstate = si_->allocState();

	unsigned int uiNumIterations = 0;
	while (sol->solution == NULL && ptc == false)
	{
		++uiNumIterations;

		/* sample random state (with goal biasing) */
		if (goal_s && rng.uniform01() < goalBias_ && goal_s->canSample())
			goal_s->sampleGoal(rstate);
		else
			samplerArray_[tid]->sampleUniform(rstate);

		/* find closest state in the tree */
		nnLock_.lock();
		Motion *nmotion = nn_->nearest(rmotion);
		nnLock_.unlock();
		base::State *dstate = rstate;

		/* find state to add */
		double d = si_->distance(nmotion->state, rstate);
		if (d > maxDistance_)
		{
			si_->getStateSpace()->interpolate(nmotion->state, rstate, maxDistance_ / d, xstate);
			dstate = xstate;
		}

		if (si_->checkMotion(nmotion->state, dstate))
		{
			/* create a motion */
			Motion *motion = new Motion(si_);
			si_->copyState(motion->state, dstate);
			motion->parent = nmotion;

			nnLock_.lock();
			nn_->add(motion);
			nnLock_.unlock();

			double dist = 0.0;
			bool solved = goal->isSatisfied(motion->state, &dist);
			if (solved)
			{
				sol->lock.lock();
				sol->approxdif = dist;
				sol->solution = motion;
				sol->lock.unlock();
				break;
			}
			if (dist < sol->approxdif)
			{
				sol->lock.lock();
				if (dist < sol->approxdif)
				{
					sol->approxdif = dist;
					sol->approxsol = motion;
				}
				sol->lock.unlock();
			}
		}
	}

	si_->freeState(xstate);
	if (rmotion->state)
		si_->freeState(rmotion->state);
	delete rmotion;

	sol->iterations[tid] = uiNumIterations;
}